from operator import itemgetter

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDataBaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

from model import get_model

#sqlalchemy,初始化MySql数据库连接
HOST = '192.168.110.18'
PORT = '3306'
DATABASE = 'yg-llm-demo'
USER = 'root'
PASSWORD = 'root'

MYSQL_URI = 'mysql+pymysql://{user}:{password}@{host}:{port}/{db}?charset=utf8'.format(
    user=USER, password=PASSWORD, host=HOST, port=PORT, db=DATABASE)

# 1连接MySQL数据库
db = SQLDatabase.from_uri(MYSQL_URI)

#测试数据库连接
#获取每个表名
# print(db.get_usable_table_names())
# #执行命令
# print(db.run("SELECT * FROM users;"))
# exit()

#2调用大模型
zhipuai=get_model()

# 3直接使用大模型和数据库整合
#根据用户输入的问题--生成SQL语句
test_chain = create_sql_query_chain(zhipuai, db)
# res = test_chain.invoke({"question":"请问users表下的数据有多少条？"})
# print(res)

#4提示词模板(LLM2=query+llm生成sql+sql检索)
answer_prompt=PromptTemplate.from_template(
    """给定以下用户问题、SQL语句和SQL执行后的结果。回答用户问题。
    Question:{question}
    SQL Query:{query}
    SQL Result:{result}
    回答：  """
    )

#5创建执行sql语句的工具
execute_sql_tool = QuerySQLDataBaseTool(db=db)

#6创建链--1.生成SQL，2.执行SQL语句
chain = (RunnablePassthrough.assign(query=test_chain).assign(result=itemgetter("query") | execute_sql_tool)
         | answer_prompt
         | zhipuai
         |StrOutputParser()
         )

res = chain.invoke({"question":"请问users表下的数据有多少条？"})
print(res)

